"""
Consciousness models and data structures for historical figure simulation.

This module defines the core data structures and models used throughout the consciousness
system, including state representations for Global Workspace Theory (GWT), Attention
Schema Theory (AST), and Integrated Information Theory (IIT).
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Tuple, Union
from enum import Enum
import numpy as np
import torch
from datetime import datetime


class ConsciousnessLevel(Enum):
    """Levels of consciousness based on IIT and phenomenological analysis."""
    UNCONSCIOUS = 0
    PRECONSCIOUS = 1
    CONSCIOUS = 2
    METACONSCIOUS = 3
    HYPERCONSCIOUS = 4


class AttentionMode(Enum):
    """Different modes of attention processing."""
    FOCUSED = "focused"
    DIFFUSE = "diffuse"
    SELECTIVE = "selective"
    DIVIDED = "divided"
    SUSTAINED = "sustained"


class ModalityType(Enum):
    """Types of sensory/cognitive modalities."""
    VISUAL = "visual"
    AUDITORY = "auditory"
    TEXTUAL = "textual"
    EMOTIONAL = "emotional"
    MEMORY = "memory"
    REASONING = "reasoning"
    LINGUISTIC = "linguistic"


@dataclass
class GlobalWorkspaceState:
    """
    Represents the global workspace state in Global Workspace Theory.
    
    The global workspace is where information from different cognitive modules
    is broadcast and integrated to create unified conscious experience.
    """
    
    # Core workspace information
    workspace_content: torch.Tensor = field(default_factory=lambda: torch.zeros(1024))
    broadcast_strength: float = 0.0
    integration_level: float = 0.0
    unity_index: float = 0.0
    
    # Module contributions to workspace
    module_contributions: Dict[str, torch.Tensor] = field(default_factory=dict)
    active_modules: List[str] = field(default_factory=list)
    
    # Temporal dynamics
    previous_state: Optional[torch.Tensor] = None
    state_history: List[torch.Tensor] = field(default_factory=list)
    
    # Information flow metrics
    information_flow: Dict[str, float] = field(default_factory=dict)
    cross_module_binding: Dict[Tuple[str, str], float] = field(default_factory=dict)


@dataclass
class AttentionState:
    """
    Represents attention state and schema in Attention Schema Theory.
    
    AST proposes that consciousness arises from the brain's model of its own
    attention processes - a schema that tracks and predicts attention allocation.
    """
    
    # Current attention focus
    attention_focus: torch.Tensor = field(default_factory=lambda: torch.zeros(512))
    attention_strength: float = 0.0
    attention_mode: AttentionMode = AttentionMode.FOCUSED
    
    # Attention allocation across modalities
    modality_weights: Dict[ModalityType, float] = field(default_factory=dict)
    resource_allocation: Dict[str, float] = field(default_factory=dict)
    
    # Meta-cognitive awareness
    attention_confidence: float = 0.0
    self_monitoring_strength: float = 0.0
    metacognitive_awareness: float = 0.0
    
    # Attention dynamics
    attention_history: List[torch.Tensor] = field(default_factory=list)
    switching_events: List[Dict[str, Any]] = field(default_factory=list)
    
    # Schema predictions
    predicted_attention: Optional[torch.Tensor] = None
    prediction_confidence: float = 0.0


@dataclass
class IITMetrics:
    """
    Metrics and measurements based on Integrated Information Theory.
    
    IIT provides mathematical frameworks for quantifying consciousness
    through measures of integrated information (Phi) and system coherence.
    """
    
    # Core IIT measures
    phi: float = 0.0  # Integrated information
    phi_max: float = 0.0  # Maximum integrated information
    phi_components: Dict[str, float] = field(default_factory=dict)
    
    # Temporal coherence
    temporal_coherence: float = 0.0  # T = corr(C(t), C(t-Δt))
    coherence_window: int = 10
    coherence_history: List[float] = field(default_factory=list)
    
    # Self-consistency
    self_consistency: float = 0.0  # S = sim(C_self(t₁), C_self(t₂))
    consistency_threshold: float = 0.8
    consistency_violations: int = 0
    
    # Information integration measures
    mutual_information: Dict[Tuple[str, str], float] = field(default_factory=dict)
    causal_strength: Dict[str, float] = field(default_factory=dict)
    
    # Phenomenological content
    phenomenological_richness: float = 0.0
    content_diversity: float = 0.0
    experiential_quality: Dict[str, float] = field(default_factory=dict)


@dataclass
class ModalityState:
    """
    Represents the state of a specific modality (visual, auditory, etc.).
    """
    
    modality_type: ModalityType
    activation: torch.Tensor = field(default_factory=lambda: torch.zeros(256))
    confidence: float = 0.0
    processing_depth: float = 0.0
    
    # Binding to other modalities
    cross_modal_bindings: Dict[ModalityType, float] = field(default_factory=dict)
    
    # Temporal dynamics
    activation_history: List[torch.Tensor] = field(default_factory=list)
    change_rate: float = 0.0


@dataclass
class ConsciousnessState:
    """
    Main consciousness state integrating all theoretical frameworks.
    
    This represents the complete conscious state of the system, integrating
    Global Workspace Theory, Attention Schema Theory, and Integrated Information
    Theory into a unified representation suitable for historical figure simulation.
    """
    
    # Timestamp and identity
    timestamp: datetime = field(default_factory=datetime.now)
    figure_identity: str = ""
    
    # Core consciousness level
    consciousness_level: ConsciousnessLevel = ConsciousnessLevel.CONSCIOUS
    consciousness_strength: float = 0.0
    
    # Theoretical framework states
    global_workspace: GlobalWorkspaceState = field(default_factory=GlobalWorkspaceState)
    attention_state: AttentionState = field(default_factory=AttentionState)
    iit_metrics: IITMetrics = field(default_factory=IITMetrics)
    
    # Multi-modal consciousness
    modality_states: Dict[ModalityType, ModalityState] = field(default_factory=dict)
    unified_consciousness: torch.Tensor = field(default_factory=lambda: torch.zeros(2048))
    
    # Fusion and integration
    fusion_weights: Dict[ModalityType, float] = field(default_factory=dict)
    integration_matrix: Optional[torch.Tensor] = None
    
    # Metacognitive structures
    self_awareness: float = 0.0
    introspective_access: float = 0.0
    reflective_consciousness: float = 0.0
    
    # Historical figure specific
    personality_influence: Dict[str, float] = field(default_factory=dict)
    historical_context_activation: float = 0.0
    cultural_framework_strength: float = 0.0
    
    # Performance metrics
    processing_latency: float = 0.0
    memory_usage: float = 0.0
    energy_consumption: float = 0.0
    
    def get_unified_vector(self) -> torch.Tensor:
        """
        Return a unified vector representation of the entire consciousness state.
        
        Returns:
            torch.Tensor: High-dimensional vector encoding the complete state
        """
        components = [
            self.global_workspace.workspace_content,
            self.attention_state.attention_focus,
            self.unified_consciousness
        ]
        
        # Add modality states
        for modality_state in self.modality_states.values():
            components.append(modality_state.activation)
        
        # Concatenate all components
        return torch.cat(components, dim=0)
    
    def get_consciousness_signature(self) -> Dict[str, float]:
        """
        Generate a signature characterizing this consciousness state.
        
        Returns:
            Dict[str, float]: Key metrics defining the consciousness signature
        """
        return {
            "consciousness_strength": self.consciousness_strength,
            "unity_index": self.global_workspace.unity_index,
            "attention_strength": self.attention_state.attention_strength,
            "phi": self.iit_metrics.phi,
            "temporal_coherence": self.iit_metrics.temporal_coherence,
            "self_consistency": self.iit_metrics.self_consistency,
            "self_awareness": self.self_awareness,
            "integration_level": self.global_workspace.integration_level,
            "metacognitive_awareness": self.attention_state.metacognitive_awareness,
            "phenomenological_richness": self.iit_metrics.phenomenological_richness
        }
    
    def update_history(self, max_history: int = 100):
        """
        Update the temporal history of consciousness states.
        
        Args:
            max_history: Maximum number of historical states to maintain
        """
        # Update global workspace history
        self.global_workspace.state_history.append(
            self.global_workspace.workspace_content.clone()
        )
        if len(self.global_workspace.state_history) > max_history:
            self.global_workspace.state_history.pop(0)
        
        # Update attention history
        self.attention_state.attention_history.append(
            self.attention_state.attention_focus.clone()
        )
        if len(self.attention_state.attention_history) > max_history:
            self.attention_state.attention_history.pop(0)
        
        # Update modality histories
        for modality_state in self.modality_states.values():
            modality_state.activation_history.append(
                modality_state.activation.clone()
            )
            if len(modality_state.activation_history) > max_history:
                modality_state.activation_history.pop(0)


@dataclass
class ConsciousnessTransition:
    """
    Represents a transition between consciousness states.
    
    This is useful for tracking how consciousness evolves over time
    and for implementing learning and adaptation mechanisms.
    """
    
    from_state: ConsciousnessState
    to_state: ConsciousnessState
    transition_vector: torch.Tensor = field(default_factory=lambda: torch.zeros(512))
    
    # Transition dynamics
    transition_strength: float = 0.0
    transition_type: str = ""
    causative_factors: Dict[str, float] = field(default_factory=dict)
    
    # Temporal information
    duration: float = 0.0
    smoothness: float = 0.0
    
    def compute_transition_metrics(self) -> Dict[str, float]:
        """
        Compute metrics characterizing this consciousness transition.
        
        Returns:
            Dict[str, float]: Metrics describing the transition
        """
        from_sig = self.from_state.get_consciousness_signature()
        to_sig = self.to_state.get_consciousness_signature()
        
        metrics = {}
        
        # Compute changes in key metrics
        for key in from_sig.keys():
            metrics[f"delta_{key}"] = to_sig[key] - from_sig[key]
            metrics[f"relative_delta_{key}"] = (
                (to_sig[key] - from_sig[key]) / max(from_sig[key], 1e-8)
            )
        
        # Overall transition magnitude
        from_vector = self.from_state.get_unified_vector()
        to_vector = self.to_state.get_unified_vector()
        
        metrics["transition_magnitude"] = torch.norm(to_vector - from_vector).item()
        metrics["transition_similarity"] = torch.cosine_similarity(
            from_vector.unsqueeze(0), to_vector.unsqueeze(0)
        ).item()
        
        return metrics